// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/util/counting_semaphore.h>

#include <chrono>
#include <condition_variable>
#include <cstdint>
#include <iostream>
#include <mutex>

#include <turbo/utility/status.h>

namespace nebula {
namespace util {

class CountingSemaphore::Impl {
 public:
  Impl(uint32_t initial_avail, double timeout_seconds)
      : num_permits_(initial_avail), timeout_seconds_(timeout_seconds) {}

  turbo::Status Acquire(uint32_t num_permits) {
    std::unique_lock<std::mutex> lk(mutex_);
    TURBO_RETURN_NOT_OK(CheckClosed());
    num_waiters_ += num_permits;
    waiter_cv_.notify_all();
    bool timed_out = !acquirer_cv_.wait_for(
        lk, std::chrono::nanoseconds(static_cast<int64_t>(timeout_seconds_ * 1e9)),
        [&] { return closed_ || num_permits <= num_permits_; });
    num_waiters_ -= num_permits;
    if (timed_out) {
      return turbo::invalid_argument_error("Timed out waiting for semaphore to release ", num_permits,
                             " permits.");
    }
    if (closed_) {
      return turbo::invalid_argument_error("Semaphore closed while acquiring");
    }
    num_permits_ -= num_permits;
    return turbo::OkStatus();
  }

  turbo::Status Release(uint32_t num_permits) {
    std::lock_guard<std::mutex> lg(mutex_);
    TURBO_RETURN_NOT_OK(CheckClosed());
    num_permits_ += num_permits;
    acquirer_cv_.notify_all();
    return turbo::OkStatus();
  }

  turbo::Status WaitForWaiters(uint32_t num_waiters) {
    std::unique_lock<std::mutex> lk(mutex_);
    TURBO_RETURN_NOT_OK(CheckClosed());
    if (waiter_cv_.wait_for(
            lk, std::chrono::nanoseconds(static_cast<int64_t>(timeout_seconds_ * 1e9)),
            [&] { return closed_ || num_waiters <= num_waiters_; })) {
      if (closed_) {
        return turbo::invalid_argument_error("Semaphore closed while waiting for waiters");
      }
      return turbo::OkStatus();
    }
    return turbo::invalid_argument_error("Timed out waiting for ", num_waiters,
                           " to start waiting on semaphore");
  }

  turbo::Status close() {
    std::lock_guard<std::mutex> lg(mutex_);
    TURBO_RETURN_NOT_OK(CheckClosed());
    closed_ = true;
    if (num_waiters_ > 0) {
      waiter_cv_.notify_all();
      acquirer_cv_.notify_all();
      return turbo::invalid_argument_error(
          "There were one or more threads waiting on a semaphore when it was closed");
    }
    return turbo::OkStatus();
  }

 private:
  turbo::Status CheckClosed() const {
    if (closed_) {
      return turbo::invalid_argument_error("Invalid operation on closed semaphore");
    }
    return turbo::OkStatus();
  }

  uint32_t num_permits_;
  double timeout_seconds_;
  uint32_t num_waiters_ = 0;
  bool closed_ = false;
  std::mutex mutex_;
  std::condition_variable acquirer_cv_;
  std::condition_variable waiter_cv_;
};

CountingSemaphore::CountingSemaphore(uint32_t initial_avail, double timeout_seconds)
    : impl_(new Impl(initial_avail, timeout_seconds)) {}

CountingSemaphore::~CountingSemaphore() = default;

turbo::Status CountingSemaphore::Acquire(uint32_t num_permits) {
  return impl_->Acquire(num_permits);
}
turbo::Status CountingSemaphore::Release(uint32_t num_permits) {
  return impl_->Release(num_permits);
}
turbo::Status CountingSemaphore::WaitForWaiters(uint32_t num_waiters) {
  return impl_->WaitForWaiters(num_waiters);
}
turbo::Status CountingSemaphore::close() { return impl_->close(); }

}  // namespace util
}  // namespace nebula
